package com.example.rocket.socket;

import org.springframework.web.socket.*;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class MyWebSocketHandler implements WebSocketHandler {

    //所有连接的集合
    private static final Map<String, WebSocketSession> SESSIONS = new ConcurrentHashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        String username = session.getAttributes().get("username").toString();
        SESSIONS.put(username, session);
        System.out.println("成功建立连接 ~ 用户名:" + username);
        sendMessage(session, "成功建立连接 ~ 用户名:" + username);
    }

    @Override
    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
        String msg = message.getPayload().toString();
        System.out.println(msg);
        String username = session.getAttributes().get("username").toString();
        sendMessage(username, "服务器收到消息收到" + username + "发送的消息，消息内容为：" + msg);
    }

    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.out.println("连接出错");
        if (session.isOpen()) {
            session.close();
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
        String username = session.getAttributes().get("username").toString();
        System.out.println(username + "的连接已关闭，status:" + closeStatus);
    }

    @Override
    public boolean supportsPartialMessages() {
        return false;
    }

    public static void sendMessage(String username, String message) {
        WebSocketSession webSocketSession = SESSIONS.get(username);
        sendMessage(webSocketSession, message);
    }

    public static void sendMessage(WebSocketSession session, String message) {
        if (session == null || !session.isOpen()) {
            return;
        }
        try {
            session.sendMessage(new TextMessage(message));
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    public static void fanoutMessage(String message) {
        SESSIONS.keySet().forEach(us -> sendMessage(us, message));
    }


}
